package com.gohlke.jme3;

import com.gohlke.lwo.LightwaveObject;
import com.gohlke.lwo.LightwaveObjectLayer;
import com.gohlke.lwo.LightwaveObjectPolygon;
import com.gohlke.lwo.LightwaveObjectPolygonType;
import com.gohlke.lwo.io.FormatException;
import com.gohlke.lwo.io.LightwaveObjectReader;
import com.jme3.asset.AssetInfo;
import com.jme3.asset.AssetLoader;
import com.jme3.asset.AssetManager;
import com.jme3.asset.ModelKey;
import com.jme3.material.Material;
import com.jme3.material.MaterialList;
import com.jme3.math.Vector2f;
import com.jme3.math.Vector3f;
import com.jme3.renderer.queue.RenderQueue;
import com.jme3.scene.Geometry;
import com.jme3.scene.Mesh;
import com.jme3.scene.Node;
import com.jme3.scene.VertexBuffer;
import com.jme3.scene.mesh.IndexBuffer;
import com.jme3.scene.mesh.IndexIntBuffer;
import com.jme3.scene.mesh.IndexShortBuffer;
import com.jme3.util.BufferUtils;
import com.jme3.util.IntMap;

import java.io.IOException;
import java.io.InputStream;
import java.nio.FloatBuffer;
import java.nio.IntBuffer;
import java.nio.ShortBuffer;
import java.util.*;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * <p>TODO: This class needs a description!</p>
 * <p>Date: 3/18/13</p>
 * <p>Time: 5:02 PM</p>
 *
 * @author jgohlke
 */
public class LightwaveObjectLoader implements AssetLoader
{
	private static final Logger logger = Logger.getLogger( LightwaveObjectLoader.class.getName() );

	protected ModelKey key = null;
	protected AssetManager assetManager = null;
	protected String objName = null;

	protected MaterialList matList = null;

	protected final HashMap<Vertex, Integer> vertIndexMap = new HashMap<Vertex, Integer>( 100 );
	protected final IntMap<Vertex> indexVertMap = new IntMap<Vertex>( 100 );
	protected final HashMap<Vertex, List<Face>> vertexToFaces = new HashMap<Vertex, List<Face>>( 100 );
	protected final ArrayList<Face> faces = new ArrayList<Face>();
	protected final Set<Vertex> vertices = new HashSet<Vertex>();

	protected int curIndex = 0;
	protected int geomIndex = 0;

	protected LightwaveObject lwObject;

	public LightwaveObjectLoader()
	{
		this.lwObject = null;
		this.key = null;
		this.assetManager = null;
		this.matList = null;
	}

	public void reset()
	{
		vertIndexMap.clear();
		indexVertMap.clear();
		faces.clear();
		vertexToFaces.clear();

		lwObject = null;
		matList = null;
		curIndex = 0;
		geomIndex = 0;
	}

	protected static class Vertex
	{
		Vector3f v;
		Vector2f vt;
		Vector3f vn;

		int meshIndex;

		@Override
		public boolean equals( Object o )
		{
			if( this == o ) return true;
			if( o == null || getClass() != o.getClass() ) return false;

			Vertex vertex = (Vertex)o;

			if( meshIndex != vertex.meshIndex ) return false;
			if( v != null ? !v.equals( vertex.v ) : vertex.v != null ) return false;

			return true;
		}

		@Override
		public int hashCode()
		{
			int result = v != null ? v.hashCode() : 0;
			result = 31 * result + meshIndex;
			return result;
		}
	}

	protected static class Face
	{
		Vertex[] verticies = null;
		LightwaveObjectPolygon lwoPolygon = null;
		Vector3f normal = null;

		public Vector3f getNormal()
		{
			if( normal == null )
				return calcNormal();
			else
				return normal;
		}

		public Vector3f calcNormal()
		{
			Vector3f v1 = verticies[ 0 ].v;
			Vector3f v2 = verticies[ 1 ].v;
			Vector3f v3 = verticies[ 2 ].v;

			Vector3f a = v1.subtract( v2 );
			Vector3f b = v1.subtract( v3 );

			return ( normal = a.cross( b ).normalize() );
		}

		@Override
		public boolean equals( Object o )
		{
			if( this == o ) return true;
			if( o == null || getClass() != o.getClass() ) return false;

			Face face = (Face)o;

			if( !Arrays.equals( verticies, face.verticies ) ) return false;

			return true;
		}

		@Override
		public int hashCode()
		{
			return verticies != null ? Arrays.hashCode( verticies ) : 0;
		}
	}

	protected Face[] quadToTriangle( Face f )
	{
		assert f.verticies.length == 4;

		Face[] t = new Face[]{ new Face(), new Face() };
		t[ 0 ].verticies = new Vertex[ 3 ];
		t[ 1 ].verticies = new Vertex[ 3 ];

		Vertex v0 = f.verticies[ 0 ];
		Vertex v1 = f.verticies[ 1 ];
		Vertex v2 = f.verticies[ 2 ];
		Vertex v3 = f.verticies[ 3 ];

		// find the pair of vertices that are closest to each over
		// v0 and v2
		// OR
		// v1 and v3
		float d1 = v0.v.distanceSquared( v2.v );
		float d2 = v1.v.distanceSquared( v3.v );
		if( d1 < d2 )
		{
			// put an edge in v0, v2
			t[ 0 ].verticies[ 0 ] = v0;
			t[ 0 ].verticies[ 1 ] = v1;
			t[ 0 ].verticies[ 2 ] = v3;

			t[ 0 ].verticies[ 0 ].vn = v0.vn;
			t[ 0 ].verticies[ 1 ].vn = v1.vn;
			t[ 0 ].verticies[ 2 ].vn = v3.vn;

			t[ 1 ].verticies[ 0 ] = v1;
			t[ 1 ].verticies[ 1 ] = v2;
			t[ 1 ].verticies[ 2 ] = v3;

			t[ 1 ].verticies[ 0 ].vn = v1.vn;
			t[ 1 ].verticies[ 1 ].vn = v2.vn;
			t[ 1 ].verticies[ 2 ].vn = v3.vn;
		}
		else
		{
			// put an edge in v1, v3
			t[ 0 ].verticies[ 0 ] = v0;
			t[ 0 ].verticies[ 1 ] = v1;
			t[ 0 ].verticies[ 2 ] = v2;

			t[ 0 ].verticies[ 0 ].vn = v0.vn;
			t[ 0 ].verticies[ 1 ].vn = v1.vn;
			t[ 0 ].verticies[ 2 ].vn = v2.vn;

			t[ 1 ].verticies[ 0 ] = v0;
			t[ 1 ].verticies[ 1 ] = v2;
			t[ 1 ].verticies[ 2 ] = v3;

			t[ 1 ].verticies[ 0 ].vn = v0.vn;
			t[ 1 ].verticies[ 1 ].vn = v2.vn;
			t[ 1 ].verticies[ 2 ].vn = v3.vn;
		}

		return t;
	}

	protected void findVertexIndex( Vertex vert )
	{
		Integer index = vertIndexMap.get( vert );
		if( index != null )
		{
			vert.meshIndex = index;
		}
		else
		{
			vert.meshIndex = curIndex++;
			vertIndexMap.put( vert, vert.meshIndex );
			indexVertMap.put( vert.meshIndex, vert );
		}
	}

	protected Geometry createGeometry( ArrayList<Face> faceList, String matName ) throws IOException
	{
		if( faceList.isEmpty() )
			throw new IOException( "No geometry data to generate mesh" );

		// Create mesh from the faces
		Mesh mesh = constructMesh( faceList );

		Geometry geom = new Geometry( objName + "-geom-" + ( geomIndex++ ), mesh );

		Material material = null;
		if( matName != null && matList != null )
		{
			// Get material from material list
			material = matList.get( matName );
		}
		if( material == null )
		{
			// create default material
			material = new Material( assetManager, "Common/MatDefs/Light/Lighting.j3md" );
			material.setFloat( "Shininess", 64 );
		}
		geom.setMaterial( material );
		if( material.isTransparent() )
			geom.setQueueBucket( RenderQueue.Bucket.Transparent );
		else
			geom.setQueueBucket( RenderQueue.Bucket.Opaque );

		if( material.getMaterialDef().getName().contains( "Lighting" )
				&& mesh.getFloatBuffer( VertexBuffer.Type.Normal ) == null )
		{
			logger.log( Level.WARNING, "OBJ mesh {0} doesn't contain normals! "
					+ "It might not display correctly", geom.getName() );
		}

		return geom;
	}

	protected Mesh constructMesh( ArrayList<Face> faceList )
	{
		Mesh m = new Mesh();
		m.setMode( Mesh.Mode.Triangles );

		boolean hasTexCoord = false;
		boolean hasNormals = false;

		ArrayList<Face> newFaces = new ArrayList<Face>( faceList.size() );
		for( Face f : faceList )
		{
			for( Vertex v : f.verticies )
			{
				findVertexIndex( v );
				if( !hasTexCoord && v.vt != null )
				{
					hasTexCoord = true;
				}
				if( !hasNormals && v.vn != null )
				{
					hasNormals = true;
				}
			}

			if( f.verticies.length == 4 )
			{
				Face[] t = quadToTriangle( f );
				newFaces.add( t[ 0 ] );
				newFaces.add( t[ 1 ] );
			}
			else
			{
				newFaces.add( f );
			}
		}

		FloatBuffer posBuf = BufferUtils.createFloatBuffer( vertIndexMap.size() * 3 );
		FloatBuffer normBuf = null;
		FloatBuffer tcBuf = null;

		if( hasNormals )
		{
			normBuf = BufferUtils.createFloatBuffer( vertIndexMap.size() * 3 );
			m.setBuffer( VertexBuffer.Type.Normal, 3, normBuf );
		}
		if( hasTexCoord )
		{
			tcBuf = BufferUtils.createFloatBuffer( vertIndexMap.size() * 2 );
			m.setBuffer( VertexBuffer.Type.TexCoord, 2, tcBuf );
		}

		IndexBuffer indexBuf = null;
		if( vertIndexMap.size() >= 65536 )
		{
			// too many verticies: use intbuffer instead of shortbuffer
			IntBuffer ib = BufferUtils.createIntBuffer( newFaces.size() * 3 );
			m.setBuffer( VertexBuffer.Type.Index, 3, ib );
			indexBuf = new IndexIntBuffer( ib );
		}
		else
		{
			ShortBuffer sb = BufferUtils.createShortBuffer( newFaces.size() * 3 );
			m.setBuffer( VertexBuffer.Type.Index, 3, sb );
			indexBuf = new IndexShortBuffer( sb );
		}

		int numFaces = newFaces.size();
		for( int i = 0; i < numFaces; i++ )
		{
			Face f = newFaces.get( i );
			if( f.verticies.length != 3 )
				continue;

			Vertex v0 = f.verticies[ 0 ];
			Vertex v1 = f.verticies[ 1 ];
			Vertex v2 = f.verticies[ 2 ];

			posBuf.position( v0.meshIndex * 3 );
			posBuf.put( v0.v.x ).put( v0.v.y ).put( v0.v.z );
			posBuf.position( v1.meshIndex * 3 );
			posBuf.put( v1.v.x ).put( v1.v.y ).put( v1.v.z );
			posBuf.position( v2.meshIndex * 3 );
			posBuf.put( v2.v.x ).put( v2.v.y ).put( v2.v.z );

			if( normBuf != null )
			{
				if( v0.vn != null )
				{
					normBuf.position( v0.meshIndex * 3 );
					normBuf.put( v0.vn.x ).put( v0.vn.y ).put( v0.vn.z );
					normBuf.position( v1.meshIndex * 3 );
					normBuf.put( v1.vn.x ).put( v1.vn.y ).put( v1.vn.z );
					normBuf.position( v2.meshIndex * 3 );
					normBuf.put( v2.vn.x ).put( v2.vn.y ).put( v2.vn.z );
				}
			}

			if( tcBuf != null )
			{
				if( v0.vt != null )
				{
					tcBuf.position( v0.meshIndex * 2 );
					tcBuf.put( v0.vt.x ).put( v0.vt.y );
					tcBuf.position( v1.meshIndex * 2 );
					tcBuf.put( v1.vt.x ).put( v1.vt.y );
					tcBuf.position( v2.meshIndex * 2 );
					tcBuf.put( v2.vt.x ).put( v2.vt.y );
				}
			}

			int index = i * 3; // current face * 3 = current index
			indexBuf.put( index, v0.meshIndex );
			indexBuf.put( index + 1, v1.meshIndex );
			indexBuf.put( index + 2, v2.meshIndex );
		}

		m.setBuffer( VertexBuffer.Type.Position, 3, posBuf );
		// index buffer and others were set on creation

		m.setStatic();
		m.updateBound();
		m.updateCounts();
		//m.setInterleaved();

		// clear data generated face statements
		// to prepare for next mesh
		vertIndexMap.clear();
		indexVertMap.clear();
		curIndex = 0;

		return m;
	}

	@Override
	public Object load( AssetInfo assetInfo ) throws IOException
	{
		reset();

		key = (ModelKey)assetInfo.getKey();
		assetManager = assetInfo.getManager();
		objName = key.getName();

		String folderName = key.getFolder();
		String ext = key.getExtension();
		objName = objName.substring( 0, objName.length() - ext.length() - 1 );
		if( folderName != null && folderName.length() > 0 )
		{
			objName = objName.substring( folderName.length() );
		}

		Node objNode = new Node( objName + "-objnode" );

		if( !( assetInfo.getKey() instanceof ModelKey ) )
			throw new IllegalArgumentException( "Model assets must be loaded using a ModelKey" );

		InputStream is = assetInfo.openStream();

		try
		{
			LightwaveObjectReader reader = new LightwaveObjectReader( is );

			lwObject = reader.getObject();
		}
		catch( FormatException e )
		{
			logger.log( Level.SEVERE, "Could not read LWO file!", e );
		}
		finally
		{
			try {
				if( is != null )
					is.close();
			} catch( IOException e ){
				//Suppress
			}
		}

		if( lwObject == null )
		{
			return null;
		}

		for( LightwaveObjectLayer layer : lwObject.layers )
		{
			for( LightwaveObjectPolygon polygon : layer.polygons )
			{
				Face face = new Face();
				face.lwoPolygon = polygon;

				face.verticies = new Vertex[ polygon.vertexIndices.length ];

				if( polygon.type != LightwaveObjectPolygonType.FACE )
				{
					logger.warning( "Skipping polygon with type " + polygon.type.name() );
					continue;
				}

				if( polygon.vertexIndices.length > 4 || polygon.vertexIndices.length <= 2 )
				{
					logger.warning( "Ignoring point or edge." );
					continue;
				}

				//logger.info( Arrays.toString( polygon.vertexIndices ) );

				for( int i = 0; i < polygon.vertexIndices.length; i++ )
				{
					Vertex vertex = new Vertex();

					vertex.v = layer.vertices.get( polygon.vertexIndices[ i ] );

					face.verticies[ i ] = vertex;

					vertices.add( vertex );

					List<Face> vertexToFacesList = vertexToFaces.get( vertex );
					if( vertexToFacesList == null )
						vertexToFaces.put( vertex, ( vertexToFacesList = new ArrayList<Face>() ) );
					vertexToFacesList.add( face );
				}
				faces.add( face );
			}

			for( Vertex vertex : vertices )
			{
				Vector3f normal = new Vector3f();
				List<Face> vertexToFacesList = vertexToFaces.get( vertex );
				for( Face face : vertexToFacesList )
				{
					normal = normal.add( face.getNormal() );
				}

				//logger.info( normal.toString() + " / " + vertexToFacesList.size() );

				vertex.vn = normal.normalize();

				for( Face face : vertexToFacesList )
				{
					for( int i = 0; i < face.verticies.length; i++ )
					{
						if( face.verticies[ i ].equals( vertex ) )
							face.verticies[ i ].vn = vertex.vn;
					}
				}
			}

			Geometry geometry = createGeometry( faces, null );
			objNode.attachChild( geometry );

			vertexToFaces.clear();
			vertices.clear();
			faces.clear();
		}

		if( objNode.getQuantity() == 1 )
			// only 1 geometry, so no need to send node
			return objNode.getChild( 0 );
		else
			return objNode;
	}
}
